package tech.mhuang.pacebox.springboot.autoconfiguration.trace.rest;

import com.alibaba.fastjson.JSON;
import io.opentracing.Scope;
import io.opentracing.Span;
import io.opentracing.Tracer;
import io.opentracing.contrib.spring.web.client.HttpHeadersCarrier;
import io.opentracing.contrib.spring.web.client.RestTemplateSpanDecorator;
import io.opentracing.contrib.spring.web.client.TracingRestTemplateInterceptor;
import io.opentracing.propagation.Format;
import io.opentracing.tag.Tags;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpRequest;
import org.springframework.http.client.ClientHttpRequestExecution;
import org.springframework.http.client.ClientHttpResponse;

import java.io.IOException;
import java.util.Collections;
import java.util.List;
import java.util.regex.Pattern;

/**
 * restTemplate埋点
 *
 * @author mhuang
 * @since 1.0.0
 */
@Slf4j
public class TraceRestTemplateInterceptor extends TracingRestTemplateInterceptor {

    private Tracer tracer;
    private Pattern skipPattern;
    private List<RestTemplateSpanDecorator> spanDecorators;

    public TraceRestTemplateInterceptor(Tracer tracer, Pattern skipPattern) {
        this.tracer = tracer;
        this.skipPattern = skipPattern;
        this.spanDecorators = Collections.singletonList(new TraceRestTemplateSpanDecorator());
    }

    protected boolean isTraced(HttpRequest httpRequest) {
        // skip URLs matching skip pattern
        // e.g. pattern is defined as '/health|/status' then URL 'http://localhost:5000/context/health' won't be traced
        if (skipPattern != null) {
            String url = httpRequest.getURI().toString();
            return skipPattern.matcher(url).matches();
        }

        return false;
    }

    @Override
    public ClientHttpResponse intercept(HttpRequest httpRequest, byte[] body,
                                        ClientHttpRequestExecution execution) throws IOException {
        if (isTraced(httpRequest)) {
            return execution.execute(httpRequest, body);
        }
        ClientHttpResponse httpResponse;

        Span span = tracer.buildSpan(httpRequest.getMethod().toString())
                .withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_CLIENT)
                .start();
        span.setTag("request.body", new String(body));
        span.setTag("request.param", httpRequest.getURI().getQuery() == null ?
                ""
                :
                httpRequest.getURI().getQuery());
        span.setTag("request.header", JSON.toJSONString(httpRequest.getHeaders()));
        tracer.inject(span.context(), Format.Builtin.HTTP_HEADERS,
                new HttpHeadersCarrier(httpRequest.getHeaders()));
        for (RestTemplateSpanDecorator spanDecorator : spanDecorators) {
            try {
                spanDecorator.onRequest(httpRequest, span);
            } catch (RuntimeException exDecorator) {
                log.error("Exception during decorating span", exDecorator);
            }
        }

        try (Scope scope = tracer.activateSpan(span)) {
            httpResponse = execution.execute(httpRequest, body);
            for (RestTemplateSpanDecorator spanDecorator : spanDecorators) {
                try {
                    spanDecorator.onResponse(httpRequest, httpResponse, span);
                } catch (RuntimeException exDecorator) {
                    log.error("Exception during decorating span", exDecorator);
                }
            }
        } catch (Exception ex) {
            for (RestTemplateSpanDecorator spanDecorator : spanDecorators) {
                try {
                    spanDecorator.onError(httpRequest, ex, span);
                } catch (RuntimeException exDecorator) {
                    log.error("Exception during decorating span", exDecorator);
                }
            }
            throw ex;
        } finally {
            span.finish();
        }

        return httpResponse;
    }
}
